/**
 * 
 */
package edu.umd.clip.lm.model.training;

import java.io.IOException;
import java.util.*;

import edu.umd.clip.lm.factors.FactorTupleDescription;
import edu.umd.clip.lm.model.Experiment;
import edu.umd.clip.lm.model.data.Context;
import edu.umd.clip.lm.model.data.ContextFuturesPair;
import edu.umd.clip.lm.model.data.ReadableTrainingData;
import edu.umd.clip.lm.model.data.TrainingDataBlock;
import edu.umd.clip.lm.model.data.TupleCountPair;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class NgramCounts {
	byte order;
	Trie theTrie;
	
	/**
	 * @param order
	 */
	public NgramCounts(byte order) {
		this.order = order;
		this.theTrie = new Trie(0);
	}

	public static class Trie implements Comparable<Trie>{
		TreeMap<Trie,Trie> subTries;
		final long word;
		int count;
		
		/**
		 * @param word
		 * @param count
		 */
		public Trie(long word, int count) {
			this.word = word;
			this.count = count;
			this.subTries = null;
		}

		/**
		 * @param word
		 */
		public Trie(long word) {
			this(word, 0);
		}
		
		public Trie getSubtrie(long w) {
			if (subTries == null) return null;
			Trie testTrie = new Trie(w);
			return subTries.get(testTrie);
		}
		
		public void addSubtrie(Trie t) {
			if (subTries == null) {
				subTries = new TreeMap<Trie,Trie>();
			}
			subTries.put(t, t);
		}
		
		public void addCount(long w, int count) {
			if (subTries == null) {
				subTries = new TreeMap<Trie, Trie>();
			}
			Trie newTrie = new Trie(w, count);
			Trie trie = subTries.get(newTrie);
			if (trie == null) {
				subTries.put(newTrie, newTrie);
			} else {
				trie.count += count;
			}
		}
		
		public Trie findTrie(long ctx[]) {
			return findTrie(ctx, 0, false);
		}

		public Trie findTrie(long ctx[], boolean createIfNotFound) {
			return findTrie(ctx, 0, createIfNotFound);
		}
		
		Trie findTrie(long ctx[], int idx, boolean createIfNotFound) {
			long w = ctx[idx];
			Trie subTrie = getSubtrie(w);
			if (subTrie == null) {
				if (createIfNotFound) {
					subTrie = new Trie(w);
					addSubtrie(subTrie);
				} else {
					return null;
				}
			}
			if (idx == ctx.length - 1) {
				return subTrie;
			}
			return subTrie.findTrie(ctx, idx+1, createIfNotFound);
		}

		// makes sure that count is the sum of subTries counts (if subTries is not null)
		public void updateCounts() {
			if (subTries == null || subTries.size() == 0) return;
			int totalCount = 0;
			for(Trie trie : subTries.keySet()) {
				trie.updateCounts();
				totalCount += trie.count;
			}
			count = totalCount;
		}
		
		private class LevelIterator implements Iterator<Trie> {
			final int level;
			Iterator<Trie> subTrieIterator;
			LevelIterator previous;
			
			/**
			 * @param level
			 */
			public LevelIterator(int level) {
				this.level = level;
				if (level < 1 || subTries == null) {
					throw new IllegalArgumentException();
				}
				this.subTrieIterator = subTries.keySet().iterator();
				this.previous = null;
			}

			/* (non-Javadoc)
			 * @see java.util.Iterator#hasNext()
			 */
			@Override
			public boolean hasNext() {
				if (level == 1) return subTrieIterator.hasNext();
				
				if (previous == null || !previous.hasNext()) {
					while(subTrieIterator.hasNext()) {
						Trie trie = subTrieIterator.next();
						if (previous == null || !previous.hasNext()) {
							previous = trie.new LevelIterator(level-1);
							if (previous.hasNext()) return true;
						}
					}
					return false;
				}
				return true;
			}

			/* (non-Javadoc)
			 * @see java.util.Iterator#next()
			 */
			@Override
			public Trie next() {
				if (level == 1) {
					return subTrieIterator.next();
				}
				if (hasNext()) return previous.next();
				
				throw new NoSuchElementException();
			}

			/* (non-Javadoc)
			 * @see java.util.Iterator#remove()
			 */
			@Override
			public void remove() {
				throw new UnsupportedOperationException();
			}
		}
		
		public Iterator<Trie> iterateLevel(int level) {
			return new LevelIterator(level);
		}
		
		/* (non-Javadoc)
		 * @see java.lang.Comparable#compareTo(java.lang.Object)
		 */
		@Override
		public int compareTo(Trie o) {
			if (word < o.word) return -1;
			if (word > o.word) return 1;
			return 0;
		}

		@Override
		public int hashCode() {
			final int prime = 31;
			int result = 1;
			result = prime * result + (int) (word ^ (word >>> 32));
			return result;
		}

		@Override
		public boolean equals(Object obj) {
			if (this == obj)
				return true;
			if (obj == null)
				return false;
			if (!(obj instanceof Trie))
				return false;
			Trie other = (Trie) obj;
			if (word != other.word)
				return false;
			return true;
		}
		
	}
	
	public void populateCounts(ReadableTrainingData data) throws IOException {
		FactorTupleDescription desc = Experiment.getInstance().getTupleDescription();
		final long overtMask = desc.getOvertFactorsMask();
		final long sentEndWord = desc.createEndTuple() & overtMask;
		
		while(data.hasNext()) {
			TrainingDataBlock block = data.next();
			
			long[] context = new long[order-1];
			for(ContextFuturesPair pair : block) {
				Context ctx = pair.getContext();
				assert(order == ctx.getOrder());
				
				for(int i=0; i<context.length; ++i) {
					context[i] = ctx.data[i] & overtMask;
				}
				
				Trie trie = theTrie.findTrie(context, true);
				for(TupleCountPair tc : pair.getFutures()) {
					long word = tc.tuple & overtMask;
					trie.addCount(word, tc.count);
					if (word == sentEndWord) {
						// add </s> </s> ... </s> ngrams, 
						// otherwise the first level of the trie (which is used as the unigram) will not contain </s>
						// similar for the bigram and other lower level distributions
						long newContext[] = new long[order-1];
						for(int i=1; i<=newContext.length; ++i) {
							System.arraycopy(context, i, newContext, 0, context.length-i);
							Arrays.fill(newContext, context.length-i, newContext.length, sentEndWord);
							Trie newTrie = theTrie.findTrie(newContext, true);
							newTrie.addCount(sentEndWord, tc.count);
						}
					}
				}
			}
		}
		
		theTrie.updateCounts();
	}

	public Trie findTrie(long[] ctx) {
		return theTrie.findTrie(ctx);
	}
	
	public Trie getCounts() {
		return theTrie;
	}

	public Iterator<Trie> iterateLevel(int level) {
		return theTrie.iterateLevel(level);
	}
}
